package skewreduce.framework.planner.partition;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;

import org.apache.hadoop.conf.Configurable;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.util.ReflectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import skewreduce.framework.planner.EaggCostModel;
import skewreduce.framework.planner.PartitionNode;
import skewreduce.framework.planner.axis.Axis;
import skewreduce.lib.ITuple;
import skewreduce.lib.Partition;


public abstract class Partitioner implements Configurable {
	protected static final Logger LOG = LoggerFactory.getLogger(Partitioner.class);
	
	protected Configuration conf;
	protected EaggCostModel model;
	protected List<PartitionStrategy> strategies = new ArrayList<PartitionStrategy>(8);
	protected boolean ignoreUnsafe;
	protected int unsafeDecisions;
	
	public static final class PartitionResult {
		final PartitionNode[] children;
		final boolean mergeSafe;
		
		public PartitionResult(PartitionNode... child) {
			this(true,child);
		}
		public PartitionResult(boolean safe,PartitionNode... child) {
			this.children = child;
			this.mergeSafe = safe;
		}
		
		public boolean isSafeForMerge() { return mergeSafe; }
		public PartitionNode[] getPartitions() { return children; }
	}

	@Override
	public Configuration getConf() {
		return conf;
	}

	@Override
	public void setConf(Configuration conf) {
		this.conf = conf;
	}
	
	protected void initialize(EaggCostModel model) throws Exception {
		this.model = model;
		
		// load partition strategies
		String[] classes = conf.getStrings("skewreduce.planner.partition.strategies");
		for ( String className : classes ) {
			Class<?> cls = null;
			try {
				cls = conf.getClassByName(className);
			} catch ( ClassNotFoundException x ) {
				cls = conf.getClassByName("skewreduce.framework.planner.partition.PartitionStrategy$"+className);
			}
			PartitionStrategy s = PartitionStrategy.class.cast(ReflectionUtils.newInstance(cls, conf));
			s.initialize(model);
			strategies.add(s);

            if ( LOG.isInfoEnabled() ) {
                LOG.info("loaded partitioner {}",s.getClass());
            }
		}
		
		ignoreUnsafe = conf.getBoolean("skewreduce.planner.partition.ignoreUnsafe", false);
		
		setup();
	}
	
	protected void setup() {}
	
	public final EaggCostModel getModel() { return model; }
	public final int getNumUnsafeDecisions() { return this.unsafeDecisions; }
	
	@SuppressWarnings("unchecked")
	public static Partitioner getInstance(EaggCostModel model,String type) throws Exception {
		Class<? extends Partitioner> cls = null;
		if ( "composite".equalsIgnoreCase(type) ) {
			cls = CompositePartitioner.class;
		} else if ( "bruteforce".equalsIgnoreCase(type) ) {
			cls = BruteForcePartitioner.class;
		} else {
			if ( type == null ) {
				type = model.getConf().get("skewreduce.planner.partition.axisSearch",BruteForcePartitioner.class.getCanonicalName());
			}
			try {
				cls = (Class<? extends Partitioner>) model.getConf().getClassByName(type);
			} catch (ClassNotFoundException e) {
				cls = null;
			}
			if ( cls == null )
				throw new IllegalArgumentException("Unsupported partitioner: "+type);
		}
		
		Partitioner p = ReflectionUtils.newInstance(cls, model.getConf());
		p.initialize(model);
		return p;
	}

	public abstract PartitionNode[] partition(PartitionNode root);
	
	private class PartitionContext {
		Partition split;
		PartitionResult result;
		double cost;
		double mergeCost;
		PartitionStrategy strategy;
		Axis axis;
		
		private PartitionContext() {
			reset();
		}
		
		public void set(PartitionNode root,Axis axis,PartitionResult res,double c,double mc,PartitionStrategy s) {
			this.split = root.getPartition().shallowCopy();
			this.result = res;
			this.axis = axis;
			this.cost = c;
			this.mergeCost = mc;
			this.strategy = s;
			
		}
		public boolean isSafe() { return ( result != null && result.isSafeForMerge() ); }
		public double getCost() { return cost; }
		public Partition getSplit() { return split; }
		public PartitionNode[] getPartitions() { return result.getPartitions(); }
		public double getMergeCost() { return mergeCost; }
		public Axis getAxis() { return axis; }
		public PartitionStrategy getStrategy() { return strategy; }
		
		public void reset() {
			cost = Double.MAX_VALUE;
			mergeCost = Double.MAX_VALUE;
			split = null;
			result = null;
		}
	}
	
	private PartitionContext context = new PartitionContext();
	
	static final class SplitResult {
		final int[] order;
		final float[] split;
		final float prob;
		SplitResult(int[] o,float[] s) {
			this(o,s,0.0f);
		}
		SplitResult(int[] o,float[] s,float p) {
			order = o;
			split = s;
			prob = p;
		}
		
		public int[] getAxesOrder() { return order; }
		public float[] getSplit() { return split; }
		public float getLeftProb() { return prob; }
		
		@Override
		public String toString() {
			return String.format("Axis=%s; Split=%s; leftProb = %f",Arrays.toString(order),Arrays.toString(split),prob);
		}
	}
	
	static final class TupleComparator implements Comparator<ITuple> {
		final int[] fieldOrder;
		
		TupleComparator(int[] order,int field) {
			fieldOrder = new int[order.length];
			fieldOrder[0] = field;
			for ( int i = 0, j = 1; i < fieldOrder.length; ++i ) {
				if ( order[i] == field ) continue;
				fieldOrder[j++] = order[i];
			}
		}

		@Override
		public int compare(ITuple o1,ITuple o2) {
			for ( int i = 0; i < fieldOrder.length; ++i ) {
				int fld = fieldOrder[i];
				float c = Math.signum(o1.getFloat(fld) - o2.getFloat(fld));
				if ( c != 0.0 ) return (int)(c);
			}
			return 0;
		}
		
		public int compare2(ITuple o1,ITuple o2) {
			for ( int i = 0; i < fieldOrder.length; ++i ) {
				int fld = fieldOrder[i];
				float c = Math.signum(o1.getFloat(fld) - o2.getFloat(fld));
				if ( c != 0.0 ) return (int)(c) * (i+1);
			}
			return 0;
		}
		
		public int[] getFieldOrder() { return fieldOrder; }
		
		public SplitResult computeSplit(List<ITuple> sample,int li,int ri) {
			ITuple l = sample.get(li);
			ITuple r = sample.get(ri);
			int diffField = Math.abs(compare2(l,r));
			
			float[] splitPoint = new float[ ( diffField == 0 ) ? fieldOrder.length : diffField ];
			
			float prob = 0.0f;
		
			if ( diffField == 0 ) {
				// the worst case. we should come up with the probability
				// then count probability
				for ( int i = 0; i < fieldOrder.length; ++i )
					splitPoint[i] = l.getFloat(fieldOrder[i]);
				// count left, count right. assign probability
				int countL = 0;
				int countR = 0;
				for ( int j = li; j >= 0 && l.equals(sample.get(j)); --j,++countL);
				for ( int j = ri; j < sample.size() && r.equals(sample.get(j)); ++j,++countR);
				
				prob = countL / (float)(countL+countR);
			} else {
				// we got different point for each ends. compute the split.
				int limit = splitPoint.length-1;
				for ( int i = 0; i < limit; ++i )
					splitPoint[i] = l.getFloat(fieldOrder[i]);
				splitPoint[limit] = ( l.getFloat(fieldOrder[limit]) + r.getFloat(fieldOrder[limit]) ) * 0.5f;
			}
			
			SplitResult res = new SplitResult(fieldOrder,splitPoint,prob);
			
			if ( LOG.isDebugEnabled() ) {
				LOG.debug("split info = {}",res);
			}
			
			return res;
		}
	}
	
	protected PartitionNode[] doPartition(PartitionNode root,Axis[] axes, int[] sortOrder) {
		context.reset();
		
		// enumerate all axes and choose the best partition
		for ( PartitionStrategy strategy : strategies ) {
			for ( Axis axis : axes ) {
				if ( axis == null ) {
					LOG.error("Axis is null");
					continue;
				}
				PartitionResult result = strategy.partition(axis, root, new TupleComparator(sortOrder,axis.getIndex()));
				PartitionNode[] children = result.getPartitions();
				if ( ! children[0].hasWorkCost() ) {
					children[0].setWorkCost(model.costWork(children[0]), model.getWork2Sec());
				}
				if ( ! children[1].hasWorkCost() ) {
					children[1].setWorkCost(model.costWork(children[1]), model.getWork2Sec());
				}
				double costM = model.costMerge(root, children);
				double c = Math.max(children[0].getWorkSecs(), children[1].getWorkSecs()) + costM * model.getMerge2Sec() + model.getScheduleDelay();
				
				if ( c < context.getCost() ) {
					boolean shouldKeep = ( ! context.isSafe() ) || result.isSafeForMerge();
					if ( shouldKeep ) {
						context.set(root, axis, result, c, costM, strategy);
						if ( LOG.isDebugEnabled() ) {
							LOG.debug("Best split so far by {} along axis {} = {} ({}/{})",new Object[] { context.getClass(), axis.getIndex(), root.getPartition(), c, costM * model.getMerge2Sec() });
						}
					}
				}
			}
		}
		
		if ( ! context.isSafe() ) {
			if ( ignoreUnsafe ) {
				++unsafeDecisions;
				if ( LOG.isInfoEnabled() ) {
					LOG.info("making unsafe decision for partition {}",root.getPartition());
				}
			} else {
				LOG.info("Can't partition further: {}",root.getPartition());
				return null;
			}
		}
		
		// now set the cost
		root.setMergeCost(context.getMergeCost(),model.getMerge2Sec());
		root.getPartition().set(context.getSplit());
		root.setChildren(context.getPartitions());
		
		if ( LOG.isDebugEnabled() ) {
			LOG.debug("Best split by {} along axis {} = {} ({}/{})",new Object[] { context.getStrategy().getClass(), context.getAxis().getIndex(), context.getSplit(), context.getCost(), root.getMergeSecs() });
		}
		
		return context.getPartitions();
	}
}
